import json
import random
import os

RANDOM_SEED = 42
random.seed(RANDOM_SEED)

input_dir = './data/'
output_file = 'processed_output.jsonl'

input_files = []
for root, _, files in os.walk(input_dir):
    for file in files:
        if file.endswith('.jsonl'):
            input_files.append(os.path.join(root, file))

prefix_count = 0
suffix_count = 0
total_count = 0

with open(output_file, 'w', encoding='utf-8') as outfile:
    for input_file in input_files:
        print(f"Processing file: {input_file}")
        with open(input_file, 'r', encoding='utf-8') as infile:
            for line in infile:

                entry = json.loads(line)

                code_lines = entry['code'].split('\n')
                total_lines = len(code_lines)

                if total_lines == 0:
                    # 如果代码行数太少，不执行mask
                    masked_lines = []
                    prefix = code_lines
                    suffix = []
                else:
                    mask_ratio = random.uniform(0.1, 0.5)
                    num_mask_lines = max(1, int(total_lines * mask_ratio))  # 至少mask一行

                    start_idx = random.randint(0, total_lines - num_mask_lines)
                    end_idx = start_idx + num_mask_lines

                    prefix = code_lines[:start_idx]
                    masked_lines = code_lines[start_idx:end_idx]
                    suffix = code_lines[end_idx:]

                    if not prefix:
                        prefix_count += 1
                    if not suffix:
                        suffix_count += 1
                    total_count += 1

                # 创建新的条目
                new_entry = {
                    'prefix': '\n'.join(prefix),
                    'mask': '\n'.join(masked_lines),
                    'suffix': '\n'.join(suffix)
                }
                # 写入到新的JSONL文件
                outfile.write(json.dumps(new_entry, ensure_ascii=False) + '\n')
# 输出统计信息
print(f"Files processed: {len(input_files)}")
print(f"Prefix count: {prefix_count}")
print(f"Suffix count: {suffix_count}")
print(f"Total entries processed: {total_count}")